module misc_utilities
    
use kind_mod
implicit none

private
public  mapfull, mapbound, sortme, gridlookup, logspace, linspace, normal_pdf, interp_1d, &
	& savings, savingv, savingm, saving3d, saving4d, saving5d, invpolicy, &
	& readings, readingv, readingm, reading3d, reading4d, reading5d, tauchen, &
	& CumulativeTransitionMatrix, CumulativeProbabilityVector

interface normal_pdf 
	module procedure normal_pdf_vec, normal_pdf_sca
end interface normal_pdf

contains


!================================================================================!
!                                                                                !
! CREATE THE CUMULATIVE VERSION OF A TRANSITION MATRIX OR IID PROBABILITY VECTOR !
!                                                                                !
!================================================================================!
! TRANSITION MATRIX: take in n x n matrix in which element (i,j) = Pr(x' = x_j | x = x_i), compute cumulative version 
subroutine CumulativeTransitionMatrix(CumTransMat, TransMat, n)
integer(ik), intent(in) :: n 
real(rk), intent(in), dimension(n,n) :: TransMat 
real(rk), intent(out), dimension(n,n) :: CumTransMat 
integer(ik) :: i

CumTransMat(:,1) = TransMat(:,1)
do i = 2,n
	CumTransMat(:, i) = CumTransMat(:, i-1) + TransMat(:, i)
end do 
end subroutine CumulativeTransitionMatrix

! IID PROBABILITY VECTOR
subroutine CumulativeProbabilityVector(CumProbVec, ProbVec, n)
integer(ik), intent(in) :: n 
real(rk), intent(in), dimension(n) :: ProbVec 
real(rk), intent(out), dimension(n) :: CumProbVec 
integer(ik) :: i

CumProbVec(1) = ProbVec(1)
do i = 2,n
	CumProbVec(i) = CumProbVec(i-1) + ProbVec(i)
end do 
end subroutine CumulativeProbabilityVector


!====================================!
!                                    !
! TAUCHEN METHOD TO DISCRETIZE AR(1) !
!                                    !
!====================================!
subroutine tauchen(rho, mu, sigma, n, m, svec, transmat)

integer, intent(in)                             :: n
double precision, intent(in)                    :: rho,mu,sigma,m
double precision, dimension(n,n), intent(out)   :: transmat
double precision, dimension(n), intent(out)     :: svec
double precision                                :: sigma_y, int_term, d, temp1, temp2
integer                                         :: i, j

! Step 0 : preliminary calculations
sigma_y = sigma / sqrt(1.0d0 - rho**2.0d0)

! Step 1 : find endpoints
call linspace(mu - m * sigma_y, mu + m * sigma_y, n, svec)
d = (svec(2) - svec(1)) / 2.0d0

! Step 2: calculate the transition matrix
do i = 1,n
    call normal_01_cdf((svec(1) + d - rho*svec(i) - (1.0d0-rho)*mu) / sigma,transmat(i,1))
do j = 2,n-1
    call normal_01_cdf((svec(j) + d - rho*svec(i) - (1.0d0-rho)*mu) / sigma,temp1)
    call normal_01_cdf((svec(j) - d - rho*svec(i) - (1.0d0-rho)*mu) / sigma,temp2)
    transmat(i,j) = temp1- temp2
end do
    call normal_01_cdf((svec(N) - d - rho*svec(i) - (1.0d0-rho)*mu) / sigma,transmat(i,n))
    transmat(i,n) = 1.0d0 - transmat(i,n)
end do

end subroutine tauchen

!============================!
!                            !
!    STANDARD NORMAL CDF     !
!                            !
!============================!
subroutine normal_01_cdf(x,cdf)

double precision, parameter :: a1 = 0.398942280444D0, &
                            &  a2 = 0.399903438504D0, &
                            &  a3 = 5.75885480458D0, &
                            &  a4 = 29.8213557808D0, &
                            &  a5 = 2.62433121679D0, &
                            &  a6 = 48.6959930692D0, &
                            &  a7 = 5.92885724438D0, &
                            &  b0 = 0.398942280385D0, &
                            &  b1 = 3.8052D-08, &
                            &  b2 = 1.00000615302D0, &
                            &  b3 = 3.98064794D-04, &
                            &  b4 = 1.98615381364D0, &
                            &  b5 = 0.151679116635D0, &
                            &  b6 = 5.29330324926D0, &
                            &  b7 = 4.8385912808D0, &
                            &  b8 = 15.1508972451D0, &
                            &  b9 = 0.742380924027D0, &
                            &  b10 = 30.789933034D0, &
                            &  b11 = 3.99019417011D0
double precision    :: q, y, x, cdf

!  |X| <= 1.28.
if ( abs ( x ) <= 1.28D+00 ) then
    y = 0.5D+00 * x * x
    q = 0.5D+00 - abs ( x ) * ( a1 - a2 * y / ( y + a3 - a4 / ( y + a5 &
        & + a6 / ( y + a7 ) ) ) )
!  1.28 < |X| <= 12.7
else if ( abs ( x ) <= 12.7D+00 ) then
    y = 0.5D+00 * x * x

    q = exp ( - y ) * b0 / ( abs ( x ) - b1 &
        & + b2 / ( abs ( x ) + b3 &
        & + b4 / ( abs ( x ) - b5 &
        & + b6 / ( abs ( x ) + b7 &
        & - b8 / ( abs ( x ) + b9 &
        & + b10 / ( abs ( x ) + b11 ) ) ) ) ) )
!  12.7 < |X|
else
    q = 0.0D+00
end if

! Take account of negative X.
if ( x < 0.0D+00 ) then
    cdf = q
else
    cdf = 1.0D+00 - q
end if

end subroutine normal_01_cdf


!=======================================!
!										!
! 1-Dimensional Interpolation on a grid !
!										!
!=======================================!
subroutine interp_1d(xi, xw, xval, nx, xgrid)
implicit none 
! inputs 
integer, intent(in) :: nx 
real*8, intent(in) :: xval
real*8, intent(in), dimension(nx) :: xgrid 
! outputs 
integer, intent(out) :: xi 
real*8, intent(out) :: xw 
! main calculations 
if (xval .lt. xgrid(1)) then 
	xi = 1
	xw = 1.0d0 
elseif (xval .gt. xgrid(nx)) then 
	xi = nx-1
	xw = 0.0d0 
else 
	xi = gridlookup(nx, xgrid, xval)
	xw = (xgrid(xi+1) - xval) / (xgrid(xi+1) - xgrid(xi))
end if 
end subroutine interp_1d  


!============!
!            !
! NORMAL PDF !
!            !
!============!
! -- scalar version
function normal_pdf_sca(x, mu, sigma)
implicit none 
real*8, intent(in) :: x, mu, sigma 
real*8 :: normal_pdf_sca 
real*8, parameter :: pi = 3.14159265359d0
normal_pdf_sca = dexp(-0.5d0 * ((x - mu) / sigma) ** 2) / (sigma * dsqrt(2.0d0 * pi))
end function normal_pdf_sca

! -- vector version 
function normal_pdf_vec(x, mu, sigma)
implicit none 
real*8, intent(in), dimension(:) :: x
real*8, intent(in) :: mu, sigma 
real*8, dimension(size(x,1)) :: normal_pdf_vec
integer :: n, i 
real*8, parameter :: pi = 3.14159265359d0
n = size(x,1)
do i = 1,n
	normal_pdf_vec(i) = dexp(-0.5d0 * ((x(i) - mu) / sigma) ** 2) / (sigma * dsqrt(2.0d0 * pi))
end do 
end function normal_pdf_vec

!===========================================!
!											!
! MAP BOUNDED PARAMETERS ONTO THE REAL LINE !
!											!
!===========================================!
subroutine mapfull(x, xlb, xub, ncal, param)
! undoes mapbound
integer, intent(in)	:: ncal
real*8, intent(in), dimension(ncal)	:: xub, param, xlb
real*8, intent(out), dimension(ncal)	:: x
integer	:: i
x = 0.0d0

! atanh undoes tanh
! dlog undoes exp
! -1 / log((1-x)/x) undoes logistic function
do i = 1,ncal
	x(i) = -dlog((xub(i) - param(i)) / (param(i) - xlb(i)))
end do

end subroutine mapfull

!===========================================!
!											!
! MAP BOUNDED PARAMETERS ONTO THE REAL LINE !
!											!
!===========================================!
subroutine mapbound(param, ncal, x, xlb, xub)
! undoes mapfull

integer, intent(in)	:: ncal
real*8, intent(in), dimension(ncal) 	:: x, xlb, xub
real*8, intent(out), dimension(ncal) :: param
integer :: i

! atanh undoes tanh
! dlog undoes exp
! -1 / log((1-x)/x) undoes logistic function
do i = 1,ncal
	param(i) = xlb(i) + (xub(i) - xlb(i)) / (1.0d0 + dexp(-x(i)))
end do

end subroutine mapbound

!===========================================!
!                                           !        
!   Sort a set of n function values, f(n)   !
!                                           !
!===========================================!
subroutine sortme(n, f, in0)

integer :: n, i, i0, i1, in0(n), fail
real*8 :: f(n), f0, f1 

do i = 1, n
    in0(i) = i
end do
fail = 1
do
	if (fail .eq. 0) then
		return
	end if
	fail = 0
	do i = 1, n-1
		i0 = in0(i)
		i1 = in0(i+1)
		f0 = f(i0)
		f1 = f(i1)

		if (f1 .lt. f0) then
			in0(i) = i1
			in0(i+1) = i0
			fail = 1
		end if
	end do
end do

end subroutine sortme

!==========================================================!
!                                                          !
! FUNCTION: gridlookup: find nearest (at or below) gridloc !
!                                                          !
!==========================================================!
pure function gridlookup(gridnum, xdist, xval)

integer :: gridnum, ixhigh, ixlow, ixplace 
real*8 :: xdist(gridnum), xval, gridlookup
intent(in) :: gridnum, xdist, xval

ixhigh = gridnum; ixlow = 1

do
    if ((ixhigh	- ixlow) .le. 1) then
	    exit
    end	if
    ixplace	= (ixhigh +	ixlow) / 2.0d0
    if (xdist(ixplace) .ge. xval) then
		ixhigh = ixplace
    else
		ixlow =	ixplace
    end	if
end	do
gridlookup = ixlow

end function gridlookup

!===============!
!               !
!	LogSpace	!
!               !
!===============!
subroutine logspace(lowerbound, upperbound, n, x)

integer:: n, j1
real*8:: lowerbound, upperbound, term0, lb, ub
real*8:: x(n), yval(n-1)
intent(in):: lowerbound, upperbound, n
intent(out):: x

! This function departs from its MATLAB 5.3 counterpart in that 
! lower and upper bounds are the actual bounds on the logarithmically
! space vector.
term0 = dlog(10.0d0)
lb = dlog(lowerbound) / term0
ub = dlog(upperbound) / term0

do j1 = 1, n-2
	yval(j1) = dble(j1)
end do

yval = ((ub - lb) / (n - 1)) * yval + lb

x(1) = lb
x(2:n-1) = yval
x(n) = ub

! Raising a scalar to a vector yields a vector.
x = 10.0d0**x

end subroutine logspace

!===============!
!               !
!	LinSpace	!
!               !
!===============!
subroutine linspace(lb, ub, gridnum, x)

integer :: gridnum, j1
real*8 :: lb, ub, x(gridnum), y(gridnum-2)

intent(in):: lb, ub, gridnum
intent(out):: x

! This subroutine generates a vector which contains a grid 
! of n equally spaced elements between loweround and upper-
! bound.  Note this version produces a row vector.

do j1 = 1, gridnum - 2
	y(j1) = dble(j1)
end do

! Note that I am multiplying an n-2x1 vector, Y, by a scalar 
! (which implies multiplying each element of the vector by 
! this scalar) and then adding another scalar (which implies 
! adding this scalar to each element of the vector).

y = ((ub - lb)/(gridnum-1))*y + lb

x(1) = lb
x(2:gridnum-1) = dble(y)
x(gridnum) = ub

end subroutine linspace








!====================!
!                    !
! SAVING AND READING !
!                    !
!====================!
! SCALARS
subroutine savings(x, xlabel, save_dir)

real*8, intent(in) :: x
character(len = *), intent(in) :: xlabel
character(len = *), intent(in) :: save_dir


open (unit = 1, file = trim(save_dir)//trim(xlabel), status = 'replace')
	write (unit = 1, fmt = *) x
close (unit = 1)

end subroutine savings

subroutine readings(x, xlabel, dir)
real*8, intent(out) :: x
character(len = *), intent(in) :: xlabel
character(len = *), intent(in) :: dir

open (unit = 1, file = trim(dir)//trim(xlabel), status = 'old', action = 'read')
	read(1, * ) x
close (unit = 1)

end subroutine readings


! 1D ARRAYS (VECTORS)
subroutine savingv(x, xlabel, save_dir)
! this subroutine saves an artibtrarily sized (n x 1) vector as a file called xlabel
real*8, intent(in), dimension(:) :: x
character(len = *), intent(in) :: xlabel
character(len = *), intent(in) :: save_dir
integer :: i, n

n = size(x)
open (unit = 1, file = trim(save_dir)//trim(xlabel), status = 'replace')
	do i = 1,n
		write (unit = 1, fmt = *) x(i)
	end do
close (unit = 1)

end subroutine savingv


subroutine readingv(x, n1, xlabel, dir)
integer, intent(in) :: n1 
real*8, intent(out), dimension(n1) :: x
character(len = *), intent(in) :: xlabel, dir
integer :: i
open (unit = 1, file = trim(dir)//trim(xlabel), status = 'old', action = 'read')
	read(1, * ) (x(i), i=1,n1)
close (unit = 1)

end subroutine readingv


subroutine readingvInt(x, n1, xlabel, dir)
integer, intent(in) :: n1 
integer, intent(out), dimension(n1) :: x
character(len = *), intent(in) :: xlabel, dir
integer :: i
open (unit = 1, file = trim(dir)//trim(xlabel), status = 'old', action = 'read')
	read(1, * ) (x(i), i=1,n1)
close (unit = 1)

end subroutine readingvInt



! 2D ARRAYS (MATRICES)
subroutine savingm(x, xlabel, save_dir)
real*8, intent(in), dimension(:,:) :: x
character(len = *), intent(in) :: xlabel
character(len = *), intent(in) :: save_dir

integer :: i, j, n, p
integer, dimension(2) :: np

! main
np = shape(x)
n = np(1)
p = np(2)

open (unit = 1, file = trim(save_dir)//trim(xlabel), status = 'replace')
	do j = 1,p
	do i = 1,n
		write (unit = 1, fmt = *) x(i,j)
	end do
	end do
close (unit = 1)
end subroutine savingm


subroutine readingm(x, n1, n2, xlabel, dir)
integer, intent(in) :: n1, n2
real*8, intent(out), dimension(n1, n2) :: x
character(len = *), intent(in) :: xlabel, dir
integer :: i, n
real*8, allocatable, dimension(:) :: xIn

n = n1 * n2
allocate(xIn(n))
open (unit = 1, file = trim(dir)//trim(xlabel), status = 'old', action = 'read')
	read(1, * ) (xIn(i), i=1,n)
close (unit = 1)
x = reshape(xIn, (/ n1, n2/))
deallocate(xIn)

end subroutine readingm


! 3D ARRAYS 
subroutine saving3d(x, xlabel, save_dir)
real*8, intent(in), dimension(:,:,:) :: x
character(len = *), intent(in) :: xlabel
character(len = *), intent(in) :: save_dir
integer :: i, j, h, n, p, q
integer, dimension(3) :: np

np = shape(x)
n = np(1)
p = np(2)
q = np(3)

open (unit = 1, file = trim(save_dir)//trim(xlabel), status = 'replace')
	do h = 1,q
	do j = 1,p
	do i = 1,n
		write (unit = 1, fmt = *) x(i,j,h)
	end do
	end do
	end do
close (unit = 1)
end subroutine saving3d


subroutine reading3d(x, n1, n2, n3, xlabel, dir)
integer, intent(in) :: n1, n2, n3
real*8, intent(out), dimension(n1, n2, n3) :: x
character(len = *), intent(in) :: xlabel, dir
integer :: i, n
real*8, allocatable, dimension(:) :: xIn

n = n1 * n2 * n3
allocate(xIn(n))
open (unit = 1, file = trim(dir)//trim(xlabel), status = 'old', action = 'read')
	read(1, * ) (xIn(i), i=1,n)
close (unit = 1)
x = reshape(xIn, (/ n1, n2, n3/))
deallocate(xIn)

end subroutine reading3d


! 4D ARRAYS 
subroutine saving4d(x, xlabel, save_dir)

real*8, intent(in), dimension(:,:,:,:) :: x
character(len = *), intent(in) :: xlabel
character(len = *), intent(in) :: save_dir

integer :: i, j, h,k, n, p, q
integer, dimension(4) :: np

! main
np = shape(x)
n = np(1)
p = np(2)
q = np(3)

open (unit = 1, file = trim(save_dir)//trim(xlabel), status = 'replace')
	do k = 1,np(4)
	do h = 1,q
	do j = 1,p
	do i = 1,n
		write (unit = 1, fmt = *) x(i,j,h,k)
	end do
	end do
	end do
	end do
close (unit = 1)
end subroutine saving4d


subroutine reading4d(x, n1, n2, n3, n4, xlabel, dir)
integer, intent(in) :: n1, n2, n3, n4
real*8, intent(out), dimension(n1, n2, n3, n4) :: x
character(len = *), intent(in) :: xlabel, dir
integer :: i, n
real*8, allocatable, dimension(:) :: xIn

n = n1 * n2 * n3 * n4
allocate(xIn(n))
open (unit = 1, file = trim(dir)//trim(xlabel), status = 'old', action = 'read')
	read(1, * ) (xIn(i), i=1,n)
close (unit = 1)
x = reshape(xIn, (/ n1, n2, n3, n4 /))
deallocate(xIn)

end subroutine reading4d

! 5D ARRAYS 
subroutine saving5d(x, xlabel, save_dir)
real*8, intent(in), dimension(:,:,:,:,:) :: x
character(len = *), intent(in) :: xlabel
character(len = *), intent(in) :: save_dir
integer :: i, j, h,k, n, p, q,l
integer, dimension(5) :: np

! main
np = shape(x)
n = np(1)
p = np(2)
q = np(3)

open (unit = 1, file = trim(save_dir)//trim(xlabel), status = 'replace')
	do l = 1,np(5)
	do k = 1,np(4)
	do h = 1,q
	do j = 1,p
	do i = 1,n
		write (unit = 1, fmt = *) x(i,j,h,k,l)
	end do
	end do
	end do
	end do
	end do 
close (unit = 1)
end subroutine saving5d

subroutine reading5d(x, n1, n2, n3, n4, n5, xlabel, dir)
integer, intent(in) :: n1, n2, n3, n4, n5 
real*8, intent(out), dimension(n1, n2, n3, n4, n5) :: x
character(len = *), intent(in) :: xlabel, dir
integer :: i, n
real*8, allocatable, dimension(:) :: xIn

n = n1 * n2 * n3 * n4 * n5
allocate(xIn(n))
open (unit = 1, file = trim(dir)//trim(xlabel), status = 'old', action = 'read')
	read(1, * ) (xIn(i), i=1,n)
close (unit = 1)
x = reshape(xIn, (/ n1, n2, n3, n4, n5 /))
deallocate(xIn)

end subroutine reading5d


!================================!
!                                !
! MULTIVARIATE LINEAR REGRESSION !
!								 !
!================================!
subroutine invpolicy(t, nx, nz, ivec, y, x, bval, statistics)

! Subroutine to compute linear multivariate regression with an indicator 
! variable to sort separate equations.  

integer, intent(in) :: t, nz, nx
integer, intent(in), dimension(t) :: ivec
real*8, intent(in), dimension(t)	:: y
real*8, intent(in), dimension(t, nx) :: x 
real*8, intent(inout), dimension(nz, 5) :: statistics
real*8, intent(inout), dimension(nz, nx+1) :: bval

integer :: iz, info, it, i, ind
integer, dimension(nz) :: places
real*8, dimension(nx+1, nx+1) :: xtx
real*8, dimension(nx+1) :: xty, p
real*8:: ssr, ssd, r2, errormin, errormax, n	
real*8, allocatable:: xs(:,:,:), ys(:,:), estimates(:), residuals(:), deviations(:) ! xsi(:,:), ysi(:)

! sort the data
! each section (in the depth axis) of Xs and each column of Ys contains 
! the data for a distinct value of ivec arranged consecutively 
 
allocate(ys(t, nz), xs(t, nx+1, nz))
ys = 0.0d0
xs = 0.0d0

places = 0

do iz = 1, nz
	do it = 1, t		
		if (ivec(it) .eq. iz) then
			places(iz) = places(iz) + 1
			xs(places(iz),1,iz) = 1.0d0			! Add constant term			
			xs(places(iz),2:nx+1,iz) = x(it,1:nx)	! Expanatory variables
			ys(places(iz),iz) = y(it)	 			! Dependent variable	
		end if	
	end do
end do

! regressions
do iz = 1, nz	! a separate regression for each value of ivec.
	
	xtx = matmul(transpose(xs(1:places(iz),1:nx+1,iz)), xs(1:places(iz), 1:nx+1, iz))
	xty = matmul(transpose(xs(1:places(iz),1:nx+1,iz)), ys(1:places(iz), iz))


    ! Table 4-3, page 4-161, (340 of 1029 in mklman51.pdf)
    ! call DGESV(Nx+1, nrhs, XTX, lda, ipiv, XTY, ldb, info)
    ! 
    ! INPUT 
    ! Nx+1 - order of XTX, rows of XTY
    ! nrhs - number of columns of XTY
    ! XTX, XTY the basic matrices of coefficients and constants, XTX*Beta = XTY
    ! lda - the first dimension of XTX, lda.GE.max(1,Nx+1)
    ! ldb - the first dimension of XTY, ldb.GE.max(1,Nx+1)
    !
    ! OUTPUT
    ! XTX = PLU
    ! XTY = Beta, solution matrix
    ! P - permutation matrix
    ! info - output 0 if succesful, -i if i-th parameter was illegal, i if U(i,i) 
    ! if U(i,i) = 0 and the matrix U is singular so the solution could not be 
    ! completed after LU factorisation.

    call dgesv(nx+1, 1, xtx, nx+1, p, xty, nx+1, info)

    bval(iz,1:nx+1) = xty

    allocate(estimates(places(iz)), residuals(places(iz)), deviations(places(iz)))

    n = dble(places(iz) - 2)
    estimates = matmul(xs(1:places(iz),1:nx+1,iz), bval(iz,:))
    residuals = ys(1:places(iz),iz) - estimates
    ssr = dot_product(residuals, residuals)
    ssr = ssr/n
    ssd = sum(ys(1:places(iz),iz))/dble(places(iz))
    deviations =  ys(1:places(iz),iz) - ssd
    ssd = dot_product(deviations, deviations)
    ssd = ssd/n
    r2 = 1.0d0 - ssr/ssd

    r2 = 1.0d0 - (1.0d0 - r2)*((n + 1.0d0 )/(n + 1.0d0 - dble(nx)))
    ssr = sqrt(ssr)
    errormin = minval(residuals)
    errormax = maxval(residuals)

    statistics(iz,1:5) = (/ssr, r2, errormin, errormax, n+2.0d0/)	! contains ssr, R2 and error bounds

    deallocate(estimates, residuals, deviations)
    
end do

end subroutine invpolicy

end module misc_utilities